#/usr/bin/python
# -*- coding: utf-8 -*-

## Parent class of elements:


class Element():
    '''
    The parent class "Element" is defined in order to set basic properties of elements and make
    them inherit the same framework. It is useful for a general use of elements.
    In the same idea, all the methods of the class have optional arguments that let use an instance of Element
    as a model.
    '''
    def __init__(self):
        '''
        Element constructor:
        .nodesCoo: ordered nodes coordinates
        .dofs: vector of values of d.o.f. of nodes
        .shapeFunctions: list of shape functions per node (eventually per d.o.f.)
        .integration: list of integration points and associated weight
        '''
        self.nodesCoo = []
        self.dofs = []
        self.shapeFunctions = []
        self.integration = []
        self.type = None
        self.manifolds = []
    
    def lhs_name_to_method(self,name):
        
        self.LHS_NAME_TO_METHOD = {
            "Kt": self.tangentStiffness,
            "Kt_linear": self.tangentStiffness_linear, 
            "Kt_nonlinear": self.tangentStiffness_nonlinear, 
            "Kt_geometric": self.tangentStiffness_geometric,
            "K": self.stiffness,
            "K_linear":self.stiffness_linear,
            "K_nonlinear": self.stiffness_nonlinear, 
        }
        try:
            return self.LHS_NAME_TO_METHOD[name]
        except KeyError:
            raise KeyError("Undefined LHS: "+str(name))
    
    def rhs_name_to_method(self,name):
        
        self.RHS_NAME_TO_METHOD = {
            "Fint": self.Fint,
            "Fext": self.Fext, 
            "R": self.R,  
        }
        try:
            return self.LHS_NAME_TO_METHOD[name]
        except KeyError:
            raise KeyError("Undefined LHS: "+str(name))
        
        
    def origin(self, nodesCoo=None):
        
        pass
        
    def frame(self, nodesCoo = None):
        
        return "frame_calculated"
    
    def interpolation(self, coordinates):
        
        pass
    
    def N(self,x):
        pass
    
    def dN(self,x):
        pass
        
    def stiffness(self, nodesCoo = None, dofs = None, shapeFunctions = None, integration = None):
        pass
        
    def mass(self, nodesCoo = None, dofs = None, shapeFunctions = None, integration = None):
        
        pass
    
    def tangentStiffness(self, PARTS = False, nodesCoo = None, dofs = None, shapeFunctions = None, integration = None):
        
        if nodesCoo:
            self.nodesCoo = nodesCoo
        if dofs:
            self.dofs = dofs
        if shapeFunctions:
            self.shapeFunctions = shapeFunctions
        if integration:
            self.integration = integration
        
#        TO BE CONTINUED...
#        Kt = {}
#        
#        for part in ["Kt_linear","Kt_nonlinear","Kt_geometric"]:
#            
#            Kt[part] = self.LHSPART[part](self.nodesCoo,self.shapeFunctions,self.integration)
#            try:
#                Kt["Kt"] += Kt[part]
#            except:
#                Kt["Kt"] = Kt[part]
        
#        if not part:
#            
#            Kt_l = None
#            Kt_nl = None
#            Kt_geo = None
#            
#            Kt = Kt_l + Kt_nl + Kt_geo
#            
#            return {"Kt": Kt, "Kt_linear": Kt_l, "Kt_nonlinear": Kt_nl, "Kt_geometric": Kt_geo} 
#        
#        else:
        return     
            
    def tangentStiffness_linear(self, nodesCoo = None, dofs = None, shapeFunctions = None, integration = None):        
                
        pass
    
    def tangentStiffness_nonlinear(self, nodesCoo = None, dofs = None, shapeFunctions = None, integration = None):
        
        pass
    
    def tangentStiffness_geometric(self, nodesCoo = None, dofs = None, shapeFunctions = None, integration = None):
        
        pass
    
    def Fint(self):
        pass
    
    def Fext(self):
        pass
    
    def R(self):
        pass
    
    def strains(self, nodesCoo = None, dofs = None, shapeFunctions = None, integration = None):
        
        if nodesCoo:
            self.nodesCoo = nodesCoo
        if dofs:
            self.dofs = dofs
        if shapeFunctions:
            self.shapeFunctions = shapeFunctions
        if integration:
            self.integration = integration
        
        
    
    def stresses(self, nodesCoo = None, dofs = None, shapeFunctions = None, integration = None):
        
        if nodesCoo:
            self.nodesCoo = nodesCoo
        if dofs:
            self.dofs = dofs
        if shapeFunctions:
            self.shapeFunctions = shapeFunctions
        if integration:
            self.integration = integration


    def degenerate(self, nodeIds):
        raise NotImplementedError

## Child class of class Element: 
#
#Child classes are defined for 1D to 3D elements. Each of them may has its own interpolation enriched by external
#functions. 
#The initial interpolation order (linear, quadratic or cubic) is identified in the element class name, according to the following
#format: "Element_type" + "number_of_nodes"
#
# Linearly interpolated elements   
class Rod2(Element):
    
    def __init__(self):
        Element.__init__(self)
        self.type = 'R2'
    
    def origin(self, nodesCoo=None):
        if nodesCoo:
            return (nodesCoo[0]+nodesCoo[1])/2
        else:
            return (self.nodesCoo[0]+self.nodesCoo[1])/2
        
    def degenerate(self,nodeIds):
        return ([],[])
    
class Beam2(Element):
    
    def __init__(self):
        
        self.type = 'B2'
 
    def degenerate(self,nodeIds):
        return ([],[])
    
   
class Quad4(Element):
    """ 
        4-------3
        |       |
        |       |
        1-------2
    """
    def __init__(self):
        
        pass

    def degenerate(self,nodeIds):
        
        elementsNodes = []
        elementsTypes = []
        
        if None not in [nodeIds[i] for i in [0,1]]:
            elementsNodes.append([nodeIds[i] for i in [0,1]])
            elementsTypes.append("B2")
        if None not in [nodeIds[i] for i in [1,2]]:
            elementsNodes.append([nodeIds[i] for i in [1,2]])
            elementsTypes.append("B2")
        if None not in [nodeIds[i] for i in [2,3]]:
            elementsNodes.append([nodeIds[i] for i in [2,3]])
            elementsTypes.append("B2")
        if None not in [nodeIds[i] for i in [3,0]]:
            elementsNodes.append([nodeIds[i] for i in [3,0]])
            elementsTypes.append("B2")
        
        
        return (elementsNodes,elementsTypes)
    
class Tri3(Element):
    """     3
          /   \
        /      \
        1-------2
    """
    def __init__(self):
        
        pass
    
    def degenerate(self,nodeIds):
        
        elementsNodes = []
        elementsTypes = []
        
        if None not in [nodeIds[i] for i in [0,1]]:
            elementsNodes.append([nodeIds[i] for i in [0,1]])
            elementsTypes.append("B2")
        elif None not in [nodeIds[i] for i in [1,2]]:
            elementsNodes.append([nodeIds[i] for i in [1,2]])
            elementsTypes.append("B2")
        elif None not in [nodeIds[i] for i in [2,0]]:
            elementsNodes.append([nodeIds[i] for i in [2,0]])
            elementsTypes.append("B2")
        else:
            elementsNodes  = None
            elementsTypes = None
        
        return (elementsNodes,elementsTypes)
# Quadratically interpolated elements
class Beam3(Element):
    """
       1--2--3
    """
    def __init__(self):
        
        pass

    def degenerate(self,nodeIds):
        
        elementsNodes = []
        elementsTypes = []
        
        if None not in [nodeIds[i] for i in [0,1]]:
            elementsNodes.append([nodeIds[i] for i in [0,1]])
            elementsTypes.append("B2")
        if None not in [nodeIds[i] for i in [1,2]]:
            elementsNodes.append([nodeIds[i] for i in [1,2]])
            elementsTypes.append("B2")
        
        
        return (elementsNodes,elementsTypes)
    
class Quad8(Element):
    """ 
        4---7---3
        |       |
        8       6 
        |       |
        1---5---2
    """
    def __init__(self):
        
        pass

    def degenerate(self,nodeIds):
        
        elementsNodes = []
        elementsTypes = []
        
        if None not in [nodeIds[i] for i in [0,4,1]]:
            elementsNodes.append([nodeIds[i] for i in [0,4,1]])
            elementsTypes.append("B3")
        elif None not in [nodeIds[i] for i in [0,4]]:
            elementsNodes.append([nodeIds[i] for i in [0,4]])
            elementsTypes.append("B2")
        elif None not in [nodeIds[i] for i in [4,1]]:
            elementsNodes.append([nodeIds[i] for i in [4,1]])
            elementsTypes.append("B2")  
                      
        if None not in [nodeIds[i] for i in [1,5,2]]:
            elementsNodes.append([nodeIds[i] for i in [1,5,2]])
            elementsTypes.append("B3")
        elif None not in [nodeIds[i] for i in [0,4]]:
            elementsNodes.append([nodeIds[i] for i in [1,5]])
            elementsTypes.append("B2")
        elif None not in [nodeIds[i] for i in [4,1]]:
            elementsNodes.append([nodeIds[i] for i in [5,2]])
            elementsTypes.append("B2")  
            
        if None not in [nodeIds[i] for i in [2,6,3]]:
            elementsNodes.append([nodeIds[i] for i in [2,6,3]])
            elementsTypes.append("B3")
        elif None not in [nodeIds[i] for i in [2,6]]:
            elementsNodes.append([nodeIds[i] for i in [2,6]])
            elementsTypes.append("B2")
        elif None not in [nodeIds[i] for i in [6,3]]:
            elementsNodes.append([nodeIds[i] for i in [6,3]])
            elementsTypes.append("B2")  
            
        if None not in [nodeIds[i] for i in [3,7,0]]:
            elementsNodes.append([nodeIds[i] for i in [3,7,0]])
            elementsTypes.append("B3")
        elif None not in [nodeIds[i] for i in [3,7]]:
            elementsNodes.append([nodeIds[i] for i in [3,7]])
            elementsTypes.append("B2")
        elif None not in [nodeIds[i] for i in [7,0]]:
            elementsNodes.append([nodeIds[i] for i in [7,0]])
            elementsTypes.append("B2")  
        
        return (elementsNodes,elementsTypes)

class Tri6(Element):
    
    def __init__(self):
        
        pass

# Cubically interpolated elements
eltype_to_modelclass = {
    'S3': Tri3(),
    'S4': Quad4(),
    'S6': Tri6(),
    'S8': Quad8(),
    'B2': Beam2(),
    'B3': Beam3(),
    'R2': Rod2(),          
}


